import json
import os
import pygame
from data import *
from engine import *
import tile_rpg

if __name__ == "__main__":

	def mix(a,b,func=lambda a,b:a+b):
		l = len(a)
		assert l==len(b)
		return [ func(a[i],b[i]) for i in range(0,l)]

	def render_palette(src,width):
		pos=[0,0]
		num_per_row = int(width/16)
		num_tiles = len(src['dat'])
		if num_tiles%num_per_row: t_height = int((num_tiles-(num_tiles%num_per_row))/num_per_row)+1
		else: t_height = int(num_tiles/num_per_row)
		ret = pygame.Surface((num_per_row*16,t_height*16))
		width = num_per_row*16
		for tile in [[0,i*16,16,16] for i in range(0,num_tiles)]:
			ret.blit(src['img'],pos,tile)
			pos[0]+=16
			if pos[0]>=width:pos=[0,pos[1]+16]
		return ret

	def get_tile_rect(pal,tn):
		t_p_r = int(pal.get_width()/16)
		y = int(tn/t_p_r)
		x = (tn - t_p_r*y)
		return [x*16,y*16,16,16]

	def get_tile_num(pal,pos,max):
		ret = pos[1]*int(pal.get_width()/16)+pos[0]
		assert ret<max
		return ret

	class Editor(Task):
		def __init__(self,game,map):
			Task.__init__(self,game)
			viewer_width = 480
			self.viewer = tile_rpg.MapViewer(game,map,viewer_width)
			self.viewer.set_camera([240,game.surf.get_height()//2])
			self.palette = render_palette(self.viewer.map.tileset,game.res[0]-viewer_width)
			self.clear = Clear(game)
			self.threads=[self.clear,self.viewer,self.input,self.paint_palette,self.cur]
			self.tile = 0
			self.cursor_loc = 0
			self.open_map=None

		def input_resize(self):
			c = pygame.mouse.get_pos()
			if c[0] < self.viewer.get_width():
				try:
					if self.game.pressed[pygame.K_k]:
						self.viewer.map.adjust_size(1,0,self.tile)
						self.viewer.paint_map()
					elif self.game.pressed[pygame.K_l]:
						self.viewer.map.adjust_size(-1,0,self.tile)
						self.viewer.paint_map()
					if self.game.pressed[pygame.K_j]:
						self.viewer.map.adjust_size(0,1,self.tile)
						self.viewer.paint_map()
					elif self.game.pressed[pygame.K_u]:
						self.viewer.map.adjust_size(0,-1,self.tile)
						self.viewer.paint_map()
				except:pass
				if pygame.K_RETURN in self.game.just_pressed:
					self.threads=[self.clear,self.viewer,self.input,self.paint_palette,self.cur]

		def on_terminate(self): self.viewer.map.struct = self.viewer.map.struct.__dict__

		def input(self):
			c = pygame.mouse.get_pos()
			self.cursor_loc = 0
			if c[0] < self.viewer.get_width():
				self.cursor = [int((c[0]-self.viewer.map_dest[0])/16),int((c[1]-self.viewer.map_dest[1])/16)]
				if self.game.pressed[pygame.K_UP]:self.viewer.adjust_camera(0,-16,True)
				elif self.game.pressed[pygame.K_DOWN]:self.viewer.adjust_camera(0,16,True)
				if self.game.pressed[pygame.K_LEFT]:self.viewer.adjust_camera(-16,0,True)
				elif self.game.pressed[pygame.K_RIGHT]:self.viewer.adjust_camera(16,0,True)
				if pygame.K_r in self.game.just_pressed:
					self.threads=[self.viewer,self.input_resize]
				if self.cursor[0]>=0 and self.cursor[0]<=self.viewer.map.struct.size[0]-1 and self.cursor[1]>=0 and self.cursor[1]<=self.viewer.map.struct.size[1]-1:
					self.cursor_loc = 1
					if self.game.mouse[0]:self.viewer.set_tile(self.cursor[0],self.cursor[1],self.tile)
					elif self.game.mouse[2]:self.tile=self.viewer.map.get_tile(*self.cursor)
			elif c[1]<self.palette.get_height():
				self.cursor = [int((c[0]-self.viewer.get_width())/16),int(c[1]/16)]
				if self.game.just_clicked[0]:
					try:
						self.tile=get_tile_num(self.palette,self.cursor,len(self.viewer.map.tileset['dat']))
					except:pass
				self.cursor_loc = 2

		def paint_palette(self):self.game.surf.blit(self.palette,[self.viewer.get_width(),0])

		def cur(self):
			if self.cursor_loc == 1:
				self.game.surf.blit(self.palette,mix(self.viewer.map_dest,mix(self.cursor,[16]*2,lambda a,b:a*b)),get_tile_rect(self.palette,self.tile))
			elif self.cursor_loc == 2:
				pygame.draw.rect(self.game.surf,(255,0,0),[self.viewer.get_width()+self.cursor[0]*16,self.cursor[1]*16,16,16],1)
				return
			rel_rect = get_tile_rect(self.palette,self.tile)
			rel_rect[0]+=self.viewer.get_width()
			pygame.draw.rect(self.game.surf,(255,0,0),rel_rect,2)

		def paint_tile_cur(self):pass

	class Main(Task):
		def __init__(self,game):
			Task.__init__(self,game)
			self.menu = tile_rpg.Menu(game,game.res[0],[(key,key) for key,obj in game.maps.items()])
			self.editor,self.map = None, None
			self.threads = [ self.clear, self.menu, self.check_menu]

		def clear(self):self.game.surf.fill([0]*3)

		def check_menu(self):
			if not self.menu.open:
				if self.menu.value is not None:
					self.open_map = self.menu.value
					map = tile_rpg.Map(self.game.tilesets,self.game.tileset_data,self.game.maps[self.menu.value])
					self.editor = Editor(self.game,map)
					self.threads = [ self.editor, self.input]
				self.menu.reset()

		def input(self):
			if pygame.K_ESCAPE in self.game.just_pressed:
				self.editor.terminate()
				self.game.maps[self.open_map] = self.editor.viewer.map.struct
				self.editor=None
				self.open_map=None
				self.threads=[self.clear,self.menu,self.check_menu]

		def on_terminate(self):
			if self.open_map:
				self.editor.terminate()
				self.game.maps[self.open_map] = self.editor.viewer.map.struct

	class MapEditor(Game):
		maps = load_data("maps", json.load)
		tilesets = load_data("tilesets",load_trans_image,'r+b')
		fonts = load_data("fonts", load_font,'r+b')
		tileset_data = load_data("tilesets data",json.load)
		res = (640,480)
		def __init__(self):
			Game.__init__(self, self.res)
			self.center = (int(self.res[0]/16/2),int(self.res[1]/16/2))
			self.font_table = { 'main':'fontwestern', 'menu':'fontwestern'}
			self.register_task(Main(self))

		def on_terminate(self):save_data("maps",self.maps)

	run(MapEditor())
